import abc
import concurrent.futures
import collections
import datetime
import os
import random

import numpy as np
import pandas as pd
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
import traceback
import warnings
warnings.filterwarnings('ignore')

from coin.base.datetime_util import to_datetime
from coin.proto.coin_strategy_summary_pb2 import StrategyDetailEntry, StrategyDetailSummary
from xunkemgmt_client.client.util.query_util import query_accounts

StrategyPlotMsg = collections.namedtuple('StrategyPlotMsg', [
    'strategy_name',
    'machine',
    'accounting_currency',
    'pnl',
    'pnl_from_fill',
])

SymbolPlotMsg = collections.namedtuple('SymbolPlotMsg',
                                       [
                                           'account_id',
                                           'symbol',
                                           'pnl',
                                           'position',
                                           'position_from_fill',
                                       ])

SymbolOrderPlotMsg = collections.namedtuple(
    'SymbolOrderPlotMsg',
    [
        'account_request',  # AccountRequestProto
        'symbol',  # String
        'pnl',  # List[(int, float)]
        'position',  # List[(int, float)]
        'buy_order_price',  # List[(int, float)]
        'sell_order_price',  # List[(int, float)]
        'mid_price',  # List[(int, float)]
        'start_ts',  # int
        'end_ts',  # int
    ])

StrategyOrderPlotMsg = collections.namedtuple(
    'StrategyOrderPlotMsg',
    [
        'strategy_request',  # StrategyRequestProto
        'process_info',  # ProcessInfoProto
        'symbol_plot_msg_list',  # List[SymbolOrderPlotMsg]
    ])

StrategyPlotInfo = collections.namedtuple('StrategyPlotInfo', [
    'strategy_plot_msg',  # StrategyPlotMsg
    'symbol_plot_msg_list',  # SymbolPlotMsg
])


def get_week_range(year, week):
  monday_str = '%s-W%s-1' % (year, week)
  sunday_str = '%s-W%s-0' % (year, week)
  dt_fmt = '%Y-W%W-%w'
  monday = datetime.datetime.strptime(monday_str, dt_fmt).date()
  sunday = datetime.datetime.strptime(sunday_str, dt_fmt).date()
  return monday, sunday


def plot_executor(func, plot_args_list, workers, key_func, *args):
  assert workers > 0 and isinstance(workers, int)
  if workers == 1:
    error_occured = False
    for plot_args in plot_args_list:
      key = key_func(plot_args)
      try:
        func(plot_args, *args)
      except Exception as e:
        print('Fail to plot.', key, type(e), e)
        print(traceback.format_exc())
        error_occured = True
    if error_occured:
      raise
  else:
    multiprocess_plot_executor(func, plot_args_list, workers, key_func, *args)


def multiprocess_plot_executor(func, plot_args_list, workers, key_func, *args):
  future_list = []
  assert workers > 1
  with concurrent.futures.ProcessPoolExecutor(max_workers=workers) as executor:
    for plot_args in plot_args_list:
      key = key_func(plot_args)
      future_list.append((key, executor.submit(func, plot_args, *args)))
  error_occured = False
  for key, future in future_list:
    try:
      future.result()
    except Exception as e:
      print('Fail to plot.', key, type(e), e)
      print(traceback.format_exc())
      error_occured = True
  if error_occured:
    raise


def _parse_plot_df(msg: list):
  df = pd.DataFrame(msg)
  if len(df.columns) == 2:
    return \
        [datetime.datetime.utcfromtimestamp(ts / 1e+9) for
         ts in df.iloc[:, 0].tolist()], \
        df.iloc[:, 1].tolist()
  else:
    return [], []


def _unpack_plot_info(plot_info):
  try:
    times, values = zip(*plot_info)
    times = [datetime.datetime.utcfromtimestamp(i / 10**9) for i in times]
  except Exception:
    times, values = [], []
  return times, values


def _parse_pnl_proto(strat_stat_list):
  pnl_list = []
  for strat_stat in strat_stat_list:
    time = strat_stat.stat_ts
    if strat_stat.HasField('pnl_mark'):
      value = strat_stat.pnl_mark
    elif strat_stat.HasField('pnl_mark_from_fill'):
      value = strat_stat.pnl_mark_from_fill
    else:
      value = None
    pnl_list.append((time, value))
  return pnl_list


def _parse_position_proto(account_position_list):
  position_list = []
  for account_position in account_position_list:
    assert len(account_position.each_position) == 1
    time = account_position.query_ts
    value = account_position.each_position[0].net_position
    position_list.append((time, value))
  return position_list


def _convert_proto_to_plot(strategy_detail_summary_list):
  plot_msg_list = []
  for strategy_detail_summary in strategy_detail_summary_list:
    assert isinstance(strategy_detail_summary, StrategyDetailSummary)
    strat_info = strategy_detail_summary.strategy
    process_info = strategy_detail_summary.process
    for strat_detail_entry in strategy_detail_summary.each_detail:
      pnl_list = _parse_pnl_proto(strat_detail_entry.pnl)
      pnl_from_fill_list = _parse_pnl_proto(strat_detail_entry.pnl_from_fill)
      position_list = None
      position_from_fill_list = None
      if strat_detail_entry.type == StrategyDetailEntry.SYMBOL:
        position_list = _parse_position_proto(strat_detail_entry.position)
        position_from_fill_list = _parse_position_proto(strat_detail_entry.position_from_fill)
      plot_msg_list.append([
          strat_info.strategy_name,
          process_info.machine,
          strat_detail_entry.accounting_currency,
          strat_detail_entry.type,
          strat_detail_entry.account_id,
          strat_detail_entry.symbol,
          pnl_list,
          pnl_from_fill_list,
          position_list,
          position_from_fill_list
      ])
  plot_msg = pd.DataFrame(plot_msg_list,
                          columns=[
                              'strat_name',
                              'machine',
                              'accounting_currency',
                              'type',
                              'account_id',
                              'symbol',
                              'pnl',
                              'pnl_from_fill',
                              'position',
                              'position_from_fill'
                          ])
  plot_msg.drop_duplicates(
      ['strat_name', 'machine', 'accounting_currency', 'type', 'account_id', 'symbol'],
      inplace=True)
  plot_msg.set_index(['strat_name', 'machine', 'accounting_currency'], inplace=True)
  strategy_plot_list = []
  for index in set(plot_msg.index.tolist()):
    strat_name, machine, accounting_currency = index
    strat_info = plot_msg.loc[index]
    symbol_plot_msg_list = []
    for i in range(len(strat_info)):
      if strat_info.iloc[i]['type'] == StrategyDetailEntry.STRATEGY:
        strat_plot_msg = StrategyPlotMsg(
            strategy_name=strat_name,
            machine=machine,
            accounting_currency=accounting_currency,
            pnl=strat_info.iloc[i]['pnl'],
            pnl_from_fill=strat_info.iloc[i]['pnl_from_fill'])
      elif strat_info.iloc[i]['type'] == StrategyDetailEntry.SYMBOL:
        symbol_plot_msg = SymbolPlotMsg(
            account_id=int(strat_info.iloc[i]['account_id']),
            symbol=strat_info.iloc[i]['symbol'],
            pnl=strat_info.iloc[i]['pnl'],
            position=strat_info.iloc[i]['position'],
            position_from_fill=strat_info.iloc[i]['position_from_fill'])
        symbol_plot_msg_list.append(symbol_plot_msg)
    strategy_plot_list.append(
        StrategyPlotInfo(strategy_plot_msg=strat_plot_msg,
                         symbol_plot_msg_list=symbol_plot_msg_list))
  return strategy_plot_list


class StrategyPlotter(object):
  def __init__(self, config, account_info=None):
    assert isinstance(config, StrategyPlotterConfig)
    self._config = config
    self._temp_dir = '.plot_temp'
    self._account_info = account_info or query_accounts(as_proto=True)

  @property
  def trading_date(self):
    return self._trading_date

  def _setup_plot(self, figsize_y=None):
    plt.style.use('seaborn-whitegrid')
    plt.rcParams['image.cmap'] = 'tab20c'
    plt.rcParams['lines.linewidth'] = 0.5
    plt.rcParams['figure.dpi'] = 200
    plt.rcParams['figure.figsize'] = [8, figsize_y or 3]
    plt.rcParams['font.family'] = ['sans-serif']
    plt.rcParams['font.size'] = 6
    plt.rcParams['legend.fontsize'] = 6
    plt.rcParams['xtick.labelsize'] = 6
    plt.rcParams['ytick.labelsize'] = 6
    plt.rcParams['savefig.facecolor'] = '#ffffff'

  def _get_account_owner(self, account_id):
    assert isinstance(account_id, int)
    return [x.owner for x in self._account_info if x.account_id == account_id][0]

  def _plot_balance_pos(self, strategy_plot_info, axis_interval_hour):
    assert isinstance(strategy_plot_info, StrategyPlotInfo)
    strategy_plot_msg = strategy_plot_info.strategy_plot_msg
    symbol_msg_list = strategy_plot_info.symbol_plot_msg_list
    assert isinstance(strategy_plot_msg, StrategyPlotMsg)
    strat_name = strategy_plot_msg.strategy_name
    machine = strategy_plot_msg.machine
    accounting_currency = strategy_plot_msg.accounting_currency
    subplot_rows = 2 + len(symbol_msg_list)
    fig_height = subplot_rows if subplot_rows > 6.5 else 6.5
    self._setup_plot(fig_height)
    fig = plt.figure()
    ax1 = plt.subplot2grid((subplot_rows, 1), (0, 0), rowspan=2)
    times_1, values_1 = _parse_plot_df(strategy_plot_msg.pnl)
    times_2, values_2 = _parse_plot_df(strategy_plot_msg.pnl_from_fill)
    ax1.plot(times_1, values_1, label=strat_name)
    ax1.plot(times_2, values_2, label=f'{strat_name} (from_fill)', linestyle=':')
    label_list = [strat_name]
    symbol_msg_list.sort(key=lambda x: '%s_%s' % (
        x.symbol, self._get_account_owner(x.account_id)))
    for i, symbol_msg in enumerate(symbol_msg_list):
      assert isinstance(symbol_msg, SymbolPlotMsg)
      account_id = symbol_msg.account_id
      owner = self._get_account_owner(account_id)
      symbol = symbol_msg.symbol
      symbol_label = symbol + '_' + owner
      times, values = _parse_plot_df(symbol_msg.pnl)
      ax1.plot(times, values, label=symbol_label)
      label_list.append(symbol_label)
      times_1, values_1 = _parse_plot_df(symbol_msg.position)
      times_2, values_2 = _parse_plot_df(symbol_msg.position_from_fill)
      ax2 = plt.subplot2grid((subplot_rows, 1), (2 + i, 0), sharex=ax1)
      ax2.plot(times_1, values_1, label=symbol_label, drawstyle='steps-post')
      ax2.plot(times_2, values_2, label=symbol_label, linestyle=':', drawstyle='steps-post')
      ax2.legend([symbol_label], loc='upper left')
    begin_time, end_time = self._config.get_pnl_plot_time_range()
    ax1.set_xlim((begin_time, end_time))
    if axis_interval_hour is not None:
      ax1.xaxis.set_major_locator(mdates.HourLocator(interval=axis_interval_hour))
    ax1.xaxis.set_major_formatter(mdates.DateFormatter('%m-%d %HH'))
    fig.autofmt_xdate()
    ax1.legend(loc='lower left', bbox_to_anchor=(0.0, 1.01), borderaxespad=0, ncol=3)
    fig.tight_layout(pad=0.5)
    plot_dir = self._config.get_pnl_plot_dir(strat_name, machine)
    os.makedirs(plot_dir, exist_ok=True)
    filepath = os.path.join(plot_dir, '%s.png' % ('-'.join([strat_name, accounting_currency])))
    fig.savefig(filepath, transparent=False)
    plt.cla()
    plt.close(fig)

  def _plot_unit_order(self, strategy_request, process_info, order_plot_msg):
    assert isinstance(order_plot_msg, SymbolOrderPlotMsg)
    order_plot_dict = order_plot_msg._asdict()
    begin_time = to_datetime(order_plot_dict['start_ts'])
    end_time = to_datetime(order_plot_dict['end_ts'])
    self._setup_plot(6.5)
    fig, ax1 = plt.subplots()
    ax1.set_xlim((begin_time, end_time))
    ax1.xaxis.set_major_formatter(mdates.DateFormatter('%m-%d %H:%M'))
    label = 'position'
    ax1.set_ylabel(label)
    times, values = _unpack_plot_info(order_plot_dict['position'])
    ax1.plot(times, values, label=label, drawstyle='steps-post')
    # TODO (chensili): add legend on top of plot
    # ax1.legend([label], loc='upper left')
    ax2 = ax1.twinx()
    ax2.set_ylabel('price')
    labels = ['buy_order_price', 'sell_order_price', 'mid_price']
    colors = ['green', 'red', 'goldenrod']
    for label, color in zip(labels, colors):
      alpha = 0.5 if label == 'mid_price' else 1
      times, values = _unpack_plot_info(order_plot_dict[label])
      ax2.plot(times,
               values,
               markersize=2,
               mew=0.5,
               ls='',
               marker='.',
               color=color,
               label=label,
               alpha=alpha)
    # TODO (chensili): add legend on top of plot
    # ax1.plot(order_plot_dict['pnl'])
    # ax2.legend(labels, loc='upper right')
    fig.autofmt_xdate()
    plot_dir = self._config.get_order_plot_dir(strategy_request.strategy_name,
                                               process_info.machine,
                                               order_plot_dict['account_request'],
                                               order_plot_dict['symbol'])

    os.makedirs(plot_dir, exist_ok=True)
    time_fmt = '%Y%m%dT%H%M%S'
    filename = '%s-%s.png' % (to_datetime(order_plot_dict['start_ts']).strftime(time_fmt),
                              to_datetime(order_plot_dict['end_ts']).strftime(time_fmt))
    filepath = os.path.join(plot_dir, filename)
    fig.savefig(filepath, transparent=False, dpi=200)
    plt.close(fig)

  def _plot_order(self, strategy_order_plot_msg):
    assert isinstance(strategy_order_plot_msg, StrategyOrderPlotMsg)
    for symbol_plot_msg in strategy_order_plot_msg.symbol_plot_msg_list:
      self._plot_unit_order(strategy_order_plot_msg.strategy_request,
                            strategy_order_plot_msg.process_info,
                            symbol_plot_msg)

  def do_pnl_plot(self, strategy_detail_summary_list,
                  axis_interval_hour=None, max_workers=1):
    assert max_workers > 0
    strategy_plot_list = _convert_proto_to_plot(strategy_detail_summary_list)
    key_func = lambda x: (x.strategy_plot_msg.machine, x.strategy_plot_msg.strategy_name)
    plot_executor(self._plot_balance_pos,
                  strategy_plot_list,
                  max_workers,
                  key_func,
                  axis_interval_hour)

  def do_order_plot(self, order_plot_list, axis_interval_hour=0.5, max_workers=1):
    assert max_workers > 0
    key_func = lambda x: (x.process_info.machine, x.strategy_request.strategy_name)
    plot_executor(self._plot_order, order_plot_list, max_workers, key_func)

  def update(self, **kwargs):
    self._trading_date = kwargs.get('trading_date', None) or self._trading_date
    self._root_dir = kwargs.get('root_dir', None) or self._root_dir
    self._temp_dir = kwargs.get('temp_dir', None) or self._temp_dir


class StrategyPlotterConfig(object):
  def get_pnl_plot_dir(self, strategy_name, machine):
    plot_dir = self.get_pnl_plot_dir_impl(strategy_name, machine)
    plot_dir = os.path.normpath(os.path.expanduser(plot_dir))
    return plot_dir

  def get_order_plot_dir(self, strategy_name, machine, account_request, symbol):
    plot_dir = self.get_order_plot_dir_impl(strategy_name, machine, account_request, symbol)
    plot_dir = os.path.normpath(os.path.expanduser(plot_dir))
    return plot_dir

  @abc.abstractmethod
  def get_pnl_plot_dir_impl(self, strategy_name, machine):
    raise NotImplementedError()

  @abc.abstractmethod
  def get_order_plot_dir_impl(self, strategy_name, machine, account_request, symbol):
    raise NotImplementedError()

  @abc.abstractmethod
  def get_pnl_plot_time_range(self):
    raise NotImplementedError()


class DailyStrategyPlotterConfig(StrategyPlotterConfig):
  def __init__(self, start_date, end_date, root_dir):
    self._start_date = start_date
    self._end_date = end_date
    self._pnl_plot_root_dir = os.path.join(root_dir, 'pnl_plot', 'daily')
    self._order_plot_root_dir = os.path.join(root_dir, 'order_plot', 'daily')

  def get_pnl_plot_dir_impl(self, strategy_name, machine):
    return os.path.join(self._pnl_plot_root_dir,
                        self._end_date.strftime('%Y%m%d'),
                        strategy_name,
                        machine)

  def get_order_plot_dir_impl(self, strategy_name, machine, account_request, symbol):
    acct_str = '%s_%s_%s' % (
        account_request.market_type, account_request.exchange, account_request.owner)
    return os.path.join(self._order_plot_root_dir,
                        self._end_date.strftime('%Y%m%d'),
                        strategy_name,
                        machine,
                        acct_str,
                        symbol)

  def get_pnl_plot_time_range(self):
    start_time = datetime.datetime.combine(self._start_date, datetime.time.min)
    end_time = datetime.datetime.combine(self._end_date, datetime.time.max)
    return start_time, end_time


class WeeklyStrategyPlotterConfig(StrategyPlotterConfig):
  def __init__(self, year, week, root_dir):
    self._year = year
    self._week = week
    self._pnl_plot_root_dir = os.path.join(root_dir, 'pnl_plot', 'weekly')

  def get_pnl_plot_dir_impl(self, strategy_name, machine):
    return os.path.join(self._pnl_plot_root_dir,
                        '%04dW%02d' % (self._year, self._week),
                        strategy_name,
                        machine)

  def get_pnl_plot_time_range(self):
    first_day, last_day = get_week_range(self._year, self._week)
    first_day_of_year = datetime.date(self._year, 1, 1)
    last_day_of_year = datetime.date(self._year, 12, 31)
    if first_day < first_day_of_year:
      first_day = first_day_of_year
    if last_day > last_day_of_year:
      last_day = last_day_of_year
    return \
        datetime.datetime.combine(first_day, datetime.time.min), \
        datetime.datetime.combine(last_day, datetime.time.max)
